/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
package com.webank.tikv.region;

import com.google.common.collect.RangeMap;
import com.google.common.collect.TreeRangeMap;
import com.google.protobuf.ByteString;
import com.webank.tikv.common.Key;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import static com.webank.tikv.common.Key.makeRange;

public class RegionCache {

    private final Map<Long, TiRegion> regionCache;
    private final Map<Long, TiStore> storeCache;
    private final RangeMap<Key, Long> keyToRegionIdCache;

    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();

    public RegionCache() {
        regionCache = new HashMap<>();
        storeCache = new HashMap<>();
        keyToRegionIdCache = TreeRangeMap.create();
    }

    // Do not use GetRegionByID when retrying request.
    //
    //   A,B |_______|_____|
    //   A   |_____________|
    // Consider region A, B. After merge of (A, B) -> A, region ID B does not exist.

    public TiRegion getRegionByKey(ByteString key) {
        try {
            lock.readLock().lock();
            Long regionId;
            if (key.isEmpty()) {
                // if key is empty, it must be the start key.
                regionId = keyToRegionIdCache.get(Key.toRawKey(key, true));
            } else {
                regionId = keyToRegionIdCache.get(Key.toRawKey(key));
            }
            return regionId != null ? regionCache.get(regionId) : null;
        } finally {
            lock.readLock().unlock();
        }
    }

    public TiStore getStoreById(long id) {
        try {
            lock.readLock().lock();
            return storeCache.get(id);
        } finally {
            lock.readLock().unlock();
        }
    }

    public List<TiRegion> getAllRegion() {
        try {
            lock.readLock().lock();
            return Collections.unmodifiableList(new ArrayList<>(regionCache.values()));
        } finally {
            lock.readLock().unlock();
        }
    }

    public List<TiStore> getAllStore() {
        try {
            lock.readLock().lock();
            return Collections.unmodifiableList(new ArrayList<>(storeCache.values()));
        } finally {
            lock.readLock().unlock();
        }
    }

    public TiRegion putRegion(TiRegion region) {
        try {
            lock.writeLock().lock();
            TiRegion oldRegion = regionCache.get(region.getId());
            if (oldRegion != null) {
                if (oldRegion.getMeta().equals(region.getMeta())) {
                    return oldRegion;
                }
                keyToRegionIdCache.remove(makeRange(oldRegion.getStartKey(), oldRegion.getEndKey()));
            }
            regionCache.put(region.getId(), region);
            keyToRegionIdCache.put(makeRange(region.getStartKey(), region.getEndKey()), region.getId());
            return region;
        } finally {
            lock.writeLock().unlock();
        }
    }

    public void updateRegion(TiRegion expected, TiRegion region) {
        try {
            lock.writeLock().lock();
            TiRegion oldRegion = regionCache.get(region.getId());
            if (oldRegion != null) {
                if (expected.getMeta().equals(oldRegion.getMeta())) {
                    return;
                }
                keyToRegionIdCache.remove(makeRange(oldRegion.getStartKey(), oldRegion.getEndKey()));
            }
            regionCache.put(region.getId(), region);
            keyToRegionIdCache.put(makeRange(region.getStartKey(), region.getEndKey()), region.getId());
        } finally {
            lock.writeLock().unlock();
        }
    }

    public boolean putStore(long id, TiStore store) {
        try {
            lock.writeLock().lock();
            TiStore oldStore = storeCache.get(id);
            if (oldStore != null) {
                if (oldStore.equals(store)) {
                    return false;
                }
                oldStore.markInvalid();
            }
            storeCache.put(id, store);
            return true;
        } finally {
            lock.writeLock().unlock();
        }
    }

    public boolean updateStore(TiStore oldStore, TiStore newStore) {
        if (!newStore.isValid()) {
            return false;
        }
        try {
            lock.writeLock().lock();
            if (oldStore == null) {
                storeCache.put(newStore.getId(), newStore);
                return true;
            }
            TiStore originStore = storeCache.get(oldStore.getId());
            if (originStore.equals(oldStore)) {
                storeCache.put(newStore.getId(), newStore);
                oldStore.markInvalid();
                return true;
            }
            return false;
        } finally {
            lock.writeLock().unlock();
        }
    }

    public void invalidRegion(TiRegion region) {
        try {
            lock.writeLock().lock();
            TiRegion oldRegion = regionCache.get(region.getId());
            if (oldRegion != null && oldRegion.getMeta().equals(region.getMeta())) {
                regionCache.remove(region.getId());
                keyToRegionIdCache.remove(makeRange(region.getStartKey(), region.getEndKey()));
            }
        } finally {
            lock.writeLock().unlock();
        }
    }

    public void invalidStore(long storeId) {
        try {
            lock.writeLock().lock();
            TiStore store = storeCache.remove(storeId);
            if (store != null) {
                store.markInvalid();
            }
        } finally {
            lock.writeLock().unlock();
        }
    }

    public void invalidRegionsForStore(TiStore store) {
        try {
            lock.writeLock().lock();
            TiStore oldStore = storeCache.get(store.getId());
            if (oldStore != store) {
                return;
            }
            for (TiRegion r : regionCache.values()) {
                if (r.getLeader().getStoreId() != store.getId()) {
                    continue;
                }
                keyToRegionIdCache.remove(makeRange(r.getStartKey(), r.getEndKey()));
                regionCache.remove(r.getId());
            }
        } finally {
            lock.writeLock().unlock();
        }
    }

    public void clearRegions() {
        try {
            lock.writeLock().lock();
            regionCache.clear();
            keyToRegionIdCache.clear();
        } finally {
            lock.writeLock().unlock();
        }
    }

    public void clearAll() {
        try {
            lock.writeLock().lock();
            regionCache.clear();
            storeCache.clear();
            keyToRegionIdCache.clear();
        } finally {
            lock.writeLock().unlock();
        }
    }
}
